--[[
  param: 1.限流的key
         2.每次加入的令牌数量
         3.加入令牌的时间间隔
         4.当前时间戳
         5.令牌桶容量
  return Boolean

--]]
local key = KEYS[1]
local tokenPer = tonumber(ARGV[1])
local timeInterval = tonumber(ARGV[2])
local now = tonumber(ARGV[3])
local capacity = tonumber(ARGV[4])
local ex = "EX"
-- redis中存放 令牌key和上次更新时间key
-- 判断令牌桶是否进行了初始化,进行初始化;否则取出时间key算出增加的令牌书和当前令牌数的和，判断是否限流；如果没有限流则重置key的过期时间。

local initExpire = math.ceil(capacity / tokenPer * timeInterval)
local timeKey = key..":time"
local tokenCnt = tonumber(redis.call("get", key))
if(tokenCnt == nil) then
    -- 为空设置令牌数量为初始容量
    redis.call("set", key, capacity, ex, initExpire)
    redis.call("set", timeKey, now, ex, initExpire)

    return true;
else
    local lastExpire = tonumber(redis.call("get", timeKey))
    --当前令牌数
    local sumCnt = math.floor(tokenCnt + (now - lastExpire) / timeInterval * tokenPer)
    if (sumCnt > 0) then
        -- 确保令牌数不超过桶容量
        sumCnt = math.min(sumCnt,capacity) - 1;
        redis.call("set", key, sumCnt, ex, initExpire)
        redis.call("set", timeKey, now, ex, initExpire)
        return true;
    end
    return false;

end
